package com.linxl.security.saml.demo.saml.sso;

import org.joda.time.DateTime;
import org.opensaml.common.SAMLException;
import org.opensaml.common.SAMLObject;
import org.opensaml.saml2.core.*;
import org.opensaml.saml2.metadata.SPSSODescriptor;
import org.opensaml.saml2.metadata.SingleLogoutService;
import org.opensaml.xml.XMLObject;
import org.opensaml.xml.validation.ValidationException;
import org.springframework.security.saml.context.SAMLMessageContext;
import org.springframework.security.saml.storage.SAMLMessageStorage;
import org.springframework.security.saml.websso.SingleLogoutProfileImpl;

import java.util.List;

import static org.springframework.security.saml.util.SAMLUtil.isDateTimeSkewValid;

/**
 * TDSingleLogoutProfileImpl
 *
 * @author : hui.chen
 * @version 1.0
 * @projectName：springboot-security-saml-demo
 * @className：TDSingleLogoutProfileImpl
 * @date 2023/1/15 14:25
 * @Copyright: 版权所有 (C) 2021 同盾科技.
 */
public class TDSingleLogoutProfileImpl extends SingleLogoutProfileImpl {
    @Override
    public void processLogoutResponse(SAMLMessageContext context) throws SAMLException, org.opensaml.xml.security.SecurityException, ValidationException {

        SAMLObject message = context.getInboundSAMLMessage();

        // Verify type
        if (!(message instanceof LogoutResponse)) {
            throw new SAMLException("Message is not of a LogoutResponse object type");
        }
        LogoutResponse response = (LogoutResponse) message;

        // Make sure request was authenticated if required, authentication is done as part of the binding processing
        if (!context.isInboundSAMLMessageAuthenticated() && context.getLocalExtendedMetadata().isRequireLogoutResponseSigned()) {
            throw new SAMLException("Logout Response object is required to be signed by the entity policy: " + context.getInboundSAMLMessageId());
        }

        // Verify issue time
        DateTime time = response.getIssueInstant();
        if (!isDateTimeSkewValid(getResponseSkew(), time)) {
            throw new SAMLException("Response issue time in LogoutResponse is either too old or with date in the future");
        }

        // Verify response to field if present, set request if correct
        // The inResponseTo field is optional, SAML 2.0 Core, 1542
        SAMLMessageStorage messageStorage = context.getMessageStorage();
        if (messageStorage != null && response.getInResponseTo() != null) {
            XMLObject xmlObject = messageStorage.retrieveMessage(response.getInResponseTo());
            if (xmlObject == null) {
                throw new SAMLException("InResponseToField in LogoutResponse doesn't correspond to sent message " + response.getInResponseTo());
            } else if (xmlObject instanceof LogoutRequest) {
                // Expected
            } else {
                throw new SAMLException("Sent request was of different type than the expected LogoutRequest " + response.getInResponseTo());
            }
        }

        // Verify destination
        if (response.getDestination() != null) {
            SPSSODescriptor localDescriptor = (SPSSODescriptor) context.getLocalEntityRoleMetadata();

            // Check if destination is correct on this SP
            List<SingleLogoutService> services = localDescriptor.getSingleLogoutServices();
            boolean found = false;
            for (SingleLogoutService service : services) {
                if (response.getDestination().equals(service.getLocation()) &&
                        context.getInboundSAMLBinding().equals(service.getBinding())) {
                    found = true;
                    break;
                }
            }
            if (!found) {
                throw new SAMLException("Destination in the LogoutResponse was not the expected value " + response.getDestination());
            }
        }

        // Verify issuer
        if (response.getIssuer() != null) {
            Issuer issuer = response.getIssuer();
            verifyIssuer(issuer, context);
        }

        // Verify status
        String statusCode = response.getStatus().getStatusCode().getValue();
        if (StatusCode.SUCCESS_URI.equals(statusCode)) {
            log.debug("Single Logout was successful");
        } else if (StatusCode.PARTIAL_LOGOUT_URI.equals(statusCode)) {
            log.debug("Single Logout was partially successful");
        } else {

            String message1 = response.getStatus().getStatusCode().getValue();
            String message2 = "N/A";
            StatusMessage status = response.getStatus().getStatusMessage();
            if (status != null) {
                message2 = status.getMessage();
            }
            log.warn("Received LogoutResponse has invalid status code: {}; message: {}", message1, message2);
        }
    }
}
